{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lab: Continuous Training with TensorFlow, PyTorch, XGBoost, and Scikit-learn Models with KubeFlow and AI Platform Pipelines\n",
    "\n",
    "In this lab we will create containerized training applications for ML models in TensorFlow, PyTorch, XGBoost, and Scikit-learn. Will will then use these images as ops in a KubeFlow pipeline and train multiple models in parallel. We will then set up recurring runs of our KubeFlow pipeline in the UI. \n",
    "\n",
    "First, we will containerize models in TF, PyTorch, XGBoost and Scikit-learn following a step-wise process for each:\n",
    "* Create the training script\n",
    "* Package training script into a Docker Image \n",
    "* Build and push training image to Google Cloud Container Registry\n",
    "\n",
    "Once we have all four training images built and pushed to the Container Registry, we will build a KubeFlow pipeline that does two things:\n",
    "* Queries BigQuery to create training/validation splits and export results as sharded CSV files in GCS\n",
    "* Launches AI Platform training jobs with our four containerized training applications, using the exported CSV data as input \n",
    "\n",
    "Finally, we will compile and deploy our pipeline. In the UI we will set up Continuous Training with recurring pipeline runs.\n",
    "\n",
    "**PRIOR TO STARTING THE LAB:** Make sure you create a new instance with AI Platform Pipelines. Once the GKE cluster is spun up, copy the endpoint because you will need it in this lab. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install needed package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade google-cloud-bigquery-storage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start by setting some global variables. Make sure you have created a bucket that is gs://PROJECT_ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REGION = 'us-central1'\n",
    "PROJECT_ID = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_ID[0]\n",
    "BUCKET = 'gs://' + PROJECT_ID "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, create a BigQuery dataset. We will then query a public BigQuery dataset to populate a table in this dataset. This is census data. We will use age, workclass, education, occupation, and hours per week to predict income bracket. Note: We also grab functional_weight in our query. We do not use this feature in our models, however we use it to hash on when creating training/validation splits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bq --location=US mk census"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bigquery\n",
    "    \n",
    "CREATE OR REPLACE TABLE census.data AS\n",
    "    \n",
    "SELECT age, workclass, education_num, occupation, hours_per_week,income_bracket,functional_weight \n",
    "FROM `bigquery-public-data.ml_datasets.census_adult_income` \n",
    "WHERE AGE IS NOT NULL\n",
    "AND workclass IS NOT NULL\n",
    "AND education_num IS NOT NULL\n",
    "AND occupation IS NOT NULL\n",
    "AND hours_per_week IS NOT NULL\n",
    "AND income_bracket IS NOT NULL \n",
    "AND functional_weight IS NOT NULL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Scikit-learn Training Script\n",
    "\n",
    "We will develop our first training script with Scikit-learn. We will use Pandas to read the CSV data then train a simple [SGD Classifier](https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.SGDClassifier.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir scikit_trainer_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./scikit_trainer_image/train.py\n",
    "\n",
    "\"\"\"Census Scikit classifier trainer script.\"\"\"\n",
    "\n",
    "import pickle\n",
    "import subprocess\n",
    "import sys\n",
    "import datetime\n",
    "import os\n",
    "\n",
    "import fire\n",
    "import pandas as pd\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.linear_model import SGDClassifier\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "\n",
    "def train_evaluate(training_dataset_path, validation_dataset_path,output_dir):\n",
    "    \"\"\"Trains the Census Classifier model.\"\"\"\n",
    "    \n",
    "    # Ingest data into Pandas Dataframes \n",
    "    df_train = pd.read_csv(training_dataset_path)\n",
    "    df_validation = pd.read_csv(validation_dataset_path)\n",
    "    df_train = pd.concat([df_train, df_validation])\n",
    "    \n",
    "    numeric_features = [\n",
    "        'age', 'education_num','hours_per_week'\n",
    "    ]\n",
    "    \n",
    "    categorical_features = ['workclass', 'occupation']\n",
    "    \n",
    "    # Scale numeric features, one-hot encode categorical features\n",
    "    preprocessor = ColumnTransformer(transformers=[(\n",
    "        'num', StandardScaler(),\n",
    "        numeric_features),\n",
    "        ('cat', OneHotEncoder(), categorical_features)])\n",
    "    \n",
    "    pipeline = Pipeline([('preprocessor', preprocessor),\n",
    "                         ('classifier', SGDClassifier(loss='log'))])\n",
    "    \n",
    "    num_features_type_map = {feature: 'float64' for feature in numeric_features}\n",
    "    df_train = df_train.astype(num_features_type_map)\n",
    "    df_validation = df_validation.astype(num_features_type_map)\n",
    "    \n",
    "    X_train = df_train.drop('income_bracket', axis=1)\n",
    "    y_train = df_train['income_bracket']\n",
    "    \n",
    "    # Set parameters of the model and fit\n",
    "    pipeline.set_params(classifier__alpha=0.0005, classifier__max_iter=250)\n",
    "    pipeline.fit(X_train, y_train)\n",
    "    \n",
    "    # Save the model locally\n",
    "    model_filename = 'model.pkl'\n",
    "    with open(model_filename, 'wb') as model_file:\n",
    "        pickle.dump(pipeline, model_file)\n",
    "        \n",
    "    # Copy to model to GCS \n",
    "    EXPORT_PATH = os.path.join(\n",
    "        output_dir, datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "    \n",
    "    gcs_model_path = '{}/{}'.format(EXPORT_PATH, model_filename)\n",
    "    subprocess.check_call(['gsutil', 'cp', model_filename, gcs_model_path])\n",
    "    print('Saved model in: {}'.format(gcs_model_path))\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    fire.Fire(train_evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Package Scikit-learn training script into a Docker image\n",
    "The next step is to package this training script into a Docker image. We need to be sure to list the dependencies in this Dockerfile. For the Scikit-learn model we need to ensure scikit-learn version 0.23.2 and Pandas version 1.1.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./scikit_trainer_image/Dockerfile\n",
    "\n",
    "FROM gcr.io/deeplearning-platform-release/base-cpu.py37\n",
    "RUN pip install -U fire scikit-learn==0.23.2 pandas==1.1.1\n",
    "WORKDIR /app\n",
    "COPY train.py .\n",
    "\n",
    "ENTRYPOINT [\"python\", \"train.py\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the Scikit-learn trainer image\n",
    "\n",
    "Now we will use Cloud Build to build the image and push it your project's Container Registry. Here we are using the remote cloud service to build the image, so we don't need a local installation of Docker. Note: Building and pushing the image will take a few minutes. Since we will be building and pushing 4 different images in this lab, I suggest taking a detailed look at the training scripts while you wait and make sure you understand the data ingestion/model building/training code for frameworks you develop with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCIKIT_IMAGE_NAME='scikit_trainer_image'\n",
    "SCIKIT_IMAGE_TAG='latest'\n",
    "SCIKIT_IMAGE_URI='gcr.io/{}/{}:{}'.format(PROJECT_ID, SCIKIT_IMAGE_NAME, SCIKIT_IMAGE_TAG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --tag $SCIKIT_IMAGE_URI $SCIKIT_IMAGE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create TensorFlow Training Script\n",
    "One down, three to go! Now we will develop a TensorFlow training script. We will use the tf.data API to ingest the data from CSVs then build/train a neural network with the tf.keras Functional API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir tensorflow_trainer_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./tensorflow_trainer_image/train.py\n",
    "\n",
    "\"\"\"Census Tensorflow classifier trainer script.\"\"\"\n",
    "\n",
    "import pickle\n",
    "import subprocess\n",
    "import sys\n",
    "import fire\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import datetime\n",
    "import os\n",
    "\n",
    "CSV_COLUMNS = [\"age\",\n",
    "               \"workclass\",\n",
    "               \"education_num\",\n",
    "               \"occupation\",\n",
    "               \"hours_per_week\",\n",
    "               \"income_bracket\"]\n",
    "\n",
    "# Add string name for label column\n",
    "LABEL_COLUMN = \"income_bracket\"\n",
    "\n",
    "# Set default values for each CSV column as a list of lists.\n",
    "# Treat is_male and plurality as strings.\n",
    "DEFAULTS = [[18], [\"?\"], [4], [\"?\"], [20],[\"<=50K\"]]\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    cols = tf.io.decode_csv(row_data, record_defaults=DEFAULTS)\n",
    "    feats = {\n",
    "        'age': tf.reshape(cols[0], [1,]),\n",
    "        'workclass': tf.reshape(cols[1],[1,]),\n",
    "        'education_num': tf.reshape(cols[2],[1,]),\n",
    "        'occupation': tf.reshape(cols[3],[1,]),\n",
    "        'hours_per_week': tf.reshape(cols[4],[1,]),\n",
    "        'income_bracket': cols[5]\n",
    "    }\n",
    "    label = feats.pop('income_bracket')\n",
    "    label_int = tf.case([(tf.math.equal(label,tf.constant([' <=50K'])), lambda: 0),\n",
    "                        (tf.math.equal(label,tf.constant([' >50K'])), lambda: 1)])\n",
    "    \n",
    "    return feats, label_int\n",
    "\n",
    "def load_dataset(pattern, batch_size=1, mode='eval'):\n",
    "    # Make a CSV dataset\n",
    "    filelist = tf.io.gfile.glob(pattern)\n",
    "    dataset = tf.data.TextLineDataset(filelist).skip(1)\n",
    "    dataset = dataset.map(features_and_labels)\n",
    "\n",
    "    # Shuffle and repeat for training\n",
    "    if mode == 'train':\n",
    "        dataset = dataset.shuffle(buffer_size=10*batch_size).batch(batch_size).repeat()\n",
    "    else:\n",
    "        dataset = dataset.batch(10)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "def train_evaluate(training_dataset_path, validation_dataset_path, batch_size, num_train_examples, num_evals, output_dir):\n",
    "    inputs = {\n",
    "        'age': tf.keras.layers.Input(name='age',shape=[None],dtype='int32'),\n",
    "        'workclass': tf.keras.layers.Input(name='workclass',shape=[None],dtype='string'),\n",
    "        'education_num': tf.keras.layers.Input(name='education_num',shape=[None],dtype='int32'),\n",
    "        'occupation': tf.keras.layers.Input(name='occupation',shape=[None],dtype='string'),\n",
    "        'hours_per_week': tf.keras.layers.Input(name='hours_per_week',shape=[None],dtype='int32')\n",
    "    }\n",
    "    \n",
    "    batch_size = int(batch_size)\n",
    "    num_train_examples = int(num_train_examples)\n",
    "    num_evals = int(num_evals)\n",
    "    \n",
    "    feat_cols = {\n",
    "        'age': tf.feature_column.numeric_column('age'),\n",
    "        'workclass': tf.feature_column.indicator_column(\n",
    "            tf.feature_column.categorical_column_with_hash_bucket(\n",
    "                key='workclass', hash_bucket_size=100\n",
    "            )\n",
    "        ),\n",
    "        'education_num': tf.feature_column.numeric_column('education_num'),\n",
    "        'occupation': tf.feature_column.indicator_column(\n",
    "            tf.feature_column.categorical_column_with_hash_bucket(\n",
    "                key='occupation', hash_bucket_size=100\n",
    "            )\n",
    "        ),\n",
    "        'hours_per_week': tf.feature_column.numeric_column('hours_per_week')\n",
    "    }\n",
    "    \n",
    "    dnn_inputs = tf.keras.layers.DenseFeatures(\n",
    "        feature_columns=feat_cols.values())(inputs)\n",
    "    h1 = tf.keras.layers.Dense(64, activation='relu')(dnn_inputs)\n",
    "    h2 = tf.keras.layers.Dense(128, activation='relu')(h1)\n",
    "    h3 = tf.keras.layers.Dense(64, activation='relu')(h2)\n",
    "    output = tf.keras.layers.Dense(1, activation='sigmoid')(h3)\n",
    "    \n",
    "    model = tf.keras.models.Model(inputs=inputs,outputs=output)\n",
    "    model.compile(optimizer='adam',loss='binary_crossentropy',metrics=['accuracy'])\n",
    "    \n",
    "    trainds = load_dataset(\n",
    "        pattern=training_dataset_path,\n",
    "        batch_size=batch_size,\n",
    "        mode='train')\n",
    "    \n",
    "    evalds = load_dataset(\n",
    "        pattern=validation_dataset_path,\n",
    "        mode='eval')\n",
    "    \n",
    "    \n",
    "    steps_per_epoch = num_train_examples // (batch_size * num_evals)\n",
    "    \n",
    "    history = model.fit(\n",
    "        trainds,\n",
    "        validation_data=evalds,\n",
    "        validation_steps=100,\n",
    "        epochs=num_evals,\n",
    "        steps_per_epoch=steps_per_epoch\n",
    "    )\n",
    "    \n",
    "    EXPORT_PATH = os.path.join(\n",
    "    output_dir, datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "    tf.saved_model.save(\n",
    "        obj=model, export_dir=EXPORT_PATH)  # with default serving function\n",
    "    \n",
    "    print(\"Exported trained model to {}\".format(EXPORT_PATH))\n",
    "    \n",
    "if __name__ == '__main__':\n",
    "    fire.Fire(train_evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Package TensorFlow Training Script into a Docker Image\n",
    "Note that the dependencies in this Dockerfile are different than for the Scikit-learn one. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./tensorflow_trainer_image/Dockerfile\n",
    "\n",
    "FROM gcr.io/deeplearning-platform-release/base-cpu.py37\n",
    "RUN pip install -U fire tensorflow==2.1.1\n",
    "WORKDIR /app\n",
    "COPY train.py .\n",
    "\n",
    "ENTRYPOINT [\"python\", \"train.py\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the Tensorflow Trainer Image \n",
    "Build the image and push it to your project's container registry. Again, this will take a few minutes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TF_IMAGE_NAME='tensorflow_trainer_image'\n",
    "TF_IMAGE_TAG='latest'\n",
    "TF_IMAGE_URI='gcr.io/{}/{}:{}'.format(PROJECT_ID, TF_IMAGE_NAME, TF_IMAGE_TAG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --tag $TF_IMAGE_URI $TF_IMAGE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create PyTorch Training Script\n",
    "Two down, two to go! Now we will develop a PyTorch training script. We will use Pandas DataFrames combined with PyTorch's Dataset and Dataloader to ingest the data from CSVs, build a model with torch.nn, and write a training loop to train this model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir pytorch_trainer_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./pytorch_trainer_image/train.py\n",
    "\n",
    "import os \n",
    "import subprocess\n",
    "import datetime\n",
    "import fire\n",
    "\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pandas as pd \n",
    "import numpy as np\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "class TrainData(Dataset):\n",
    "    def __init__(self, X_data, y_data):\n",
    "        self.X_data = X_data\n",
    "        self.y_data = y_data\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        return self.X_data[index], self.y_data[index]\n",
    "        \n",
    "    def __len__ (self):\n",
    "        return len(self.X_data)\n",
    "    \n",
    "class BinaryClassifier(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BinaryClassifier, self).__init__()\n",
    "        # 27 input features\n",
    "        self.h1 = nn.Linear(27, 64) \n",
    "        self.h2 = nn.Linear(64, 64)\n",
    "        self.output_layer = nn.Linear(64, 1) \n",
    "        \n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(p=0.1)\n",
    "        self.batchnorm1 = nn.BatchNorm1d(64)\n",
    "        self.batchnorm2 = nn.BatchNorm1d(64)\n",
    "        \n",
    "    def forward(self, inputs):\n",
    "        x = self.relu(self.h1(inputs))\n",
    "        x = self.batchnorm1(x)\n",
    "        x = self.relu(self.h2(x))\n",
    "        x = self.batchnorm2(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.output_layer(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "def binary_acc(y_pred, y_true):\n",
    "    \"\"\"Calculates accuracy\"\"\"\n",
    "    y_pred_tag = torch.round(torch.sigmoid(y_pred))\n",
    "\n",
    "    correct_results_sum = (y_pred_tag == y_true).sum().float()\n",
    "    acc = correct_results_sum/y_true.shape[0]\n",
    "    acc = torch.round(acc * 100)\n",
    "    \n",
    "    return acc\n",
    "\n",
    "def train_evaluate(training_dataset_path, validation_dataset_path, batch_size, num_epochs, output_dir):\n",
    "    \n",
    "    batch_size = int(batch_size)\n",
    "    num_epochs = int(num_epochs)\n",
    "    \n",
    "    # Read in train/validation data and concat \n",
    "    df_train = pd.read_csv(training_dataset_path)\n",
    "    df_validation = pd.read_csv(validation_dataset_path)\n",
    "    df = pd.concat([df_train, df_validation])\n",
    "\n",
    "    categorical_features = ['workclass', 'occupation']\n",
    "    target='income_bracket'\n",
    "\n",
    "    # One-hot encode categorical variables \n",
    "    df = pd.get_dummies(df,columns=categorical_features)\n",
    "\n",
    "    # Change label to 0 if <=50K, 1 if >50K\n",
    "    df[target] = df[target].apply(lambda x: 0 if x==' <=50K' else 1)\n",
    "\n",
    "    # Split features and labels into 2 different vars\n",
    "    X_train = df.loc[:, df.columns != target]\n",
    "    y_train = np.array(df[target])\n",
    "\n",
    "    # Normalize features \n",
    "    scaler = StandardScaler()\n",
    "    X_train = scaler.fit_transform(X_train)\n",
    "\n",
    "    # Training data\n",
    "    train_data = TrainData(torch.FloatTensor(X_train), \n",
    "                           torch.FloatTensor(y_train))\n",
    "\n",
    "    # Use torch DataLoader to feed data to model \n",
    "    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, drop_last=True)\n",
    "\n",
    "    # Instantiate model \n",
    "    model = BinaryClassifier()\n",
    "    \n",
    "    # Loss is binary crossentropy w/ logits. Must manually implement sigmoid for inference\n",
    "    criterion = nn.BCEWithLogitsLoss()\n",
    "    \n",
    "    # Adam optimizer\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "    model.train()\n",
    "    for e in range(1, num_epochs+1):\n",
    "        epoch_loss = 0\n",
    "        epoch_acc = 0\n",
    "        for X_batch, y_batch in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            y_pred = model(X_batch)\n",
    "\n",
    "            loss = criterion(y_pred, y_batch.unsqueeze(1))\n",
    "            acc = binary_acc(y_pred, y_batch.unsqueeze(1))\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "            epoch_acc += acc.item()\n",
    "\n",
    "\n",
    "        print(f'Epoch {e}: Loss = {epoch_loss/len(train_loader):.5f} | Acc = {epoch_acc/len(train_loader):.3f}')\n",
    "\n",
    "    # Save the model locally\n",
    "    model_filename='model.pt'\n",
    "    torch.save(model.state_dict(), model_filename)\n",
    "\n",
    "    EXPORT_PATH = os.path.join(\n",
    "        output_dir, datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "\n",
    "    # Copy the model to GCS\n",
    "    gcs_model_path = '{}/{}'.format(EXPORT_PATH, model_filename)\n",
    "    subprocess.check_call(['gsutil', 'cp', model_filename, gcs_model_path])\n",
    "    print('Saved model in: {}'.format(gcs_model_path))\n",
    "    \n",
    "if __name__ == '__main__':\n",
    "    fire.Fire(train_evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Package PyTorch Training Script into a Docker Image\n",
    "Note the dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./pytorch_trainer_image/Dockerfile\n",
    "\n",
    "FROM gcr.io/deeplearning-platform-release/base-cpu.py37\n",
    "RUN pip install -U fire torch==1.6.0 scikit-learn==0.23.2 pandas==1.1.1\n",
    "WORKDIR /app\n",
    "COPY train.py .\n",
    "\n",
    "ENTRYPOINT [\"python\", \"train.py\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the PyTorch Trainer Image \n",
    "Build and push the PyTorch training image to your project's Container Registry. Again, this will take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TORCH_IMAGE_NAME='pytorch_trainer_image'\n",
    "TORCH_IMAGE_TAG='latest'\n",
    "TORCH_IMAGE_URI='gcr.io/{}/{}:{}'.format(PROJECT_ID, TORCH_IMAGE_NAME, TORCH_IMAGE_TAG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --tag $TORCH_IMAGE_URI $TORCH_IMAGE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create XGBoost Training Script\n",
    "Three down, one to go! Create the final training script. This script will ingest and preprocess data with Pandas, then train a [Gradient Boosted Tree model](https://en.wikipedia.org/wiki/Gradient_boosting). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir xgboost_trainer_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./xgboost_trainer_image/train.py\n",
    "\n",
    "import os \n",
    "import subprocess\n",
    "import datetime\n",
    "import fire\n",
    "import pickle \n",
    "\n",
    "import pandas as pd \n",
    "import numpy as np\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from xgboost import XGBClassifier\n",
    "\n",
    "def train_evaluate(training_dataset_path, validation_dataset_path,max_depth,n_estimators,output_dir):\n",
    "    \n",
    "    df_train = pd.read_csv(training_dataset_path)\n",
    "    df_validation = pd.read_csv(validation_dataset_path)\n",
    "    df = pd.concat([df_train, df_validation])\n",
    "\n",
    "    categorical_features = ['workclass', 'occupation']\n",
    "    target='income_bracket'\n",
    "\n",
    "    # One-hot encode categorical variables \n",
    "    df = pd.get_dummies(df,columns=categorical_features)\n",
    "\n",
    "    # Change label to 0 if <=50K, 1 if >50K\n",
    "    df[target] = df[target].apply(lambda x: 0 if x==' <=50K' else 1)\n",
    "\n",
    "    # Split features and labels into 2 different vars\n",
    "    X_train = df.loc[:, df.columns != target]\n",
    "    y_train = np.array(df[target])\n",
    "\n",
    "    # Normalize features \n",
    "    scaler = StandardScaler()\n",
    "    X_train = scaler.fit_transform(X_train)\n",
    "    \n",
    "    grid = {\n",
    "        'max_depth': int(max_depth),\n",
    "        'n_estimators': int(n_estimators)\n",
    "    }\n",
    "    \n",
    "    model = XGBClassifier()\n",
    "    model.set_params(**grid)\n",
    "    model.fit(X_train,y_train)\n",
    "    \n",
    "    model_filename = 'xgb_model.pkl'\n",
    "    pickle.dump(model, open(model_filename, \"wb\"))\n",
    "        \n",
    "    EXPORT_PATH = os.path.join(\n",
    "        output_dir, datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "    \n",
    "    gcs_model_path = '{}/{}'.format(EXPORT_PATH, model_filename)\n",
    "    subprocess.check_call(['gsutil', 'cp', model_filename, gcs_model_path])\n",
    "    print('Saved model in: {}'.format(gcs_model_path))  \n",
    "\n",
    "if __name__ == '__main__':\n",
    "    fire.Fire(train_evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Package XGBoost Training Script into a Docker Image\n",
    "Note the dependencies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./xgboost_trainer_image/Dockerfile\n",
    "\n",
    "FROM gcr.io/deeplearning-platform-release/base-cpu.py37\n",
    "RUN pip install -U fire scikit-learn==0.23.2 pandas==1.1.1 xgboost==1.2.0\n",
    "WORKDIR /app\n",
    "COPY train.py .\n",
    "\n",
    "ENTRYPOINT [\"python\", \"train.py\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the XGBoost Trainer Image\n",
    "Build and push the XGBoost training image. This will take a few minutes (this is the last one, woohoo!) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "XGB_IMAGE_NAME='xgboost_trainer_image'\n",
    "XGB_IMAGE_TAG='latest'\n",
    "XGB_IMAGE_URI='gcr.io/{}/{}:{}'.format(PROJECT_ID, XGB_IMAGE_NAME, XGB_IMAGE_TAG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --tag $XGB_IMAGE_URI $XGB_IMAGE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Develop KubeFlow Pipeline\n",
    "Now that you have all four of your training applications as containers in your project's Container Registry, let's build a KubeFlow pipeline. \n",
    "\n",
    "The KubeFlow pipeline will have two BigQuery Ops. We will use the pre-built BigQuery Query component (no need to reinvent the wheel) to do the following: \n",
    "* Create a training split in our data and export to CSV\n",
    "* Create a validation split in our data and export to CSV\n",
    "\n",
    "The output of these BigQuery Ops will be the input data into four AI Platform Training Ops. For this we will also use a pre-built component. Each AI Platform Training Op will train one of our containerized models - Tensorflow, PyTorch, XGBoost, and Scikit-learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./pipeline/census_training_pipeline.py\n",
    "\n",
    "import os\n",
    "import kfp\n",
    "from kfp.dsl.types import GCPProjectID\n",
    "from kfp.dsl.types import GCPRegion\n",
    "from kfp.dsl.types import GCSPath\n",
    "from kfp.dsl.types import String\n",
    "from kfp.gcp import use_gcp_secret\n",
    "import kfp.components as comp\n",
    "import kfp.dsl as dsl\n",
    "import kfp.gcp as gcp\n",
    "import json\n",
    "\n",
    "# We will use environment vars to set the trainer image names and bucket name\n",
    "TF_TRAINER_IMAGE = os.getenv('TF_TRAINER_IMAGE')\n",
    "SCIKIT_TRAINER_IMAGE = os.getenv('SCIKIT_TRAINER_IMAGE')\n",
    "TORCH_TRAINER_IMAGE = os.getenv('TORCH_TRAINER_IMAGE')\n",
    "XGB_TRAINER_IMAGE = os.getenv('XGB_TRAINER_IMAGE')\n",
    "BUCKET = os.getenv('BUCKET')\n",
    "\n",
    "# Paths to export the training/validation data from bigquery\n",
    "TRAINING_OUTPUT_PATH = BUCKET + '/census/data/training.csv'\n",
    "VALIDATION_OUTPUT_PATH = BUCKET + '/census/data/validation.csv'\n",
    "\n",
    "COMPONENT_URL_SEARCH_PREFIX = 'https://raw.githubusercontent.com/kubeflow/pipelines/0.2.5/components/gcp/'\n",
    "\n",
    "# Create component factories\n",
    "component_store = kfp.components.ComponentStore(\n",
    "    local_search_paths=None, url_search_prefixes=[COMPONENT_URL_SEARCH_PREFIX])\n",
    "\n",
    "# Load BigQuery and AI Platform Training op\n",
    "bigquery_query_op = component_store.load_component('bigquery/query')\n",
    "mlengine_train_op = component_store.load_component('ml_engine/train')\n",
    "\n",
    "def get_query(dataset='training'):\n",
    "    \"\"\"Function that returns either training or validation query\"\"\"\n",
    "    if dataset=='training':\n",
    "        split = \"MOD(ABS(FARM_FINGERPRINT(CAST(functional_weight AS STRING))), 100) < 80\"\n",
    "    elif dataset=='validation':\n",
    "        split = \"\"\"MOD(ABS(FARM_FINGERPRINT(CAST(functional_weight AS STRING))), 100) >= 80 \n",
    "        AND MOD(ABS(FARM_FINGERPRINT(CAST(functional_weight AS STRING))), 100) < 90\"\"\"\n",
    "    else:\n",
    "        split = \"MOD(ABS(FARM_FINGERPRINT(CAST(functional_weight AS STRING))), 100) >= 90\"\n",
    "        \n",
    "    query = \"\"\"SELECT age, workclass, education_num, occupation, hours_per_week,income_bracket \n",
    "    FROM census.data \n",
    "    WHERE {0}\"\"\".format(split)\n",
    "    \n",
    "    return query\n",
    "\n",
    "# We will use the training/validation queries as inputs to our pipeline\n",
    "# This lets us change the training/validation datasets if we wish by simply\n",
    "# Changing the query. \n",
    "TRAIN_QUERY = get_query(dataset='training')\n",
    "VALIDATION_QUERY=get_query(dataset='validation')\n",
    "\n",
    "@dsl.pipeline(\n",
    "    name='Continuous Training with Multiple Frameworks',\n",
    "    description='Pipeline to create training/validation splits w/ BigQuery then launch multiple AI Platform Training Jobs'\n",
    ")\n",
    "def pipeline(\n",
    "    project_id,\n",
    "    train_query=TRAIN_QUERY,\n",
    "    validation_query=VALIDATION_QUERY,\n",
    "    region='us-central1'\n",
    "):\n",
    "    # Creating the training data split\n",
    "    create_training_split = bigquery_query_op(\n",
    "        query=train_query,\n",
    "        project_id=project_id,\n",
    "        output_gcs_path=TRAINING_OUTPUT_PATH\n",
    "    ).set_display_name('BQ Train Split')\n",
    "    \n",
    "    # Creating the validation data split\n",
    "    create_validation_split = bigquery_query_op(\n",
    "        query=validation_query,\n",
    "        project_id=project_id,\n",
    "        output_gcs_path=VALIDATION_OUTPUT_PATH\n",
    "    ).set_display_name('BQ Eval Split')\n",
    "    \n",
    "    # These are the output directories where our models will be saved\n",
    "    tf_output_dir = BUCKET + '/census/models/tf'\n",
    "    scikit_output_dir = BUCKET + '/census/models/scikit'\n",
    "    torch_output_dir = BUCKET + '/census/models/torch'\n",
    "    xgb_output_dir = BUCKET + '/census/models/xgb'\n",
    "    \n",
    "    # Training arguments to be passed to the TF Trainer\n",
    "    tf_args = [\n",
    "        '--training_dataset_path', create_training_split.outputs['output_gcs_path'],\n",
    "        '--validation_dataset_path', create_validation_split.outputs['output_gcs_path'],\n",
    "        '--output_dir', tf_output_dir,\n",
    "        '--batch_size', '32', \n",
    "        '--num_train_examples', '1000',\n",
    "        '--num_evals', '10'\n",
    "    ]\n",
    "    \n",
    "    # Training arguments to be passed to the Scikit-learn Trainer\n",
    "    scikit_args = [\n",
    "        '--training_dataset_path', create_training_split.outputs['output_gcs_path'],\n",
    "        '--validation_dataset_path', create_validation_split.outputs['output_gcs_path'],\n",
    "        '--output_dir', scikit_output_dir\n",
    "    ]\n",
    "    \n",
    "    # Training arguments to be passed to the PyTorch Trainer\n",
    "    torch_args = [\n",
    "        '--training_dataset_path', create_training_split.outputs['output_gcs_path'],\n",
    "        '--validation_dataset_path', create_validation_split.outputs['output_gcs_path'],\n",
    "        '--output_dir', torch_output_dir,\n",
    "        '--batch_size', '32', \n",
    "        '--num_epochs', '15',\n",
    "    ]\n",
    "    \n",
    "    # Training arguments to be passed to the XGBoost Trainer \n",
    "    xgb_args = [\n",
    "        '--training_dataset_path', create_training_split.outputs['output_gcs_path'],\n",
    "        '--validation_dataset_path', create_validation_split.outputs['output_gcs_path'],\n",
    "        '--output_dir', xgb_output_dir,\n",
    "        '--max_depth', '10', \n",
    "        '--n_estimators', '100'\n",
    "    ]\n",
    "    \n",
    "    # AI Platform Training Jobs with all 4 trainer images \n",
    "    \n",
    "    train_scikit = mlengine_train_op(\n",
    "        project_id=project_id,\n",
    "        region=region,\n",
    "        master_image_uri=SCIKIT_TRAINER_IMAGE,\n",
    "        args=scikit_args).set_display_name('Scikit Model - AI Platform Training')\n",
    "    \n",
    "    train_tf = mlengine_train_op(\n",
    "        project_id=project_id,\n",
    "        region=region,\n",
    "        master_image_uri=TF_TRAINER_IMAGE,\n",
    "        args=tf_args).set_display_name('Tensorflow Model - AI Platform Training')\n",
    "    \n",
    "    train_torch = mlengine_train_op(\n",
    "        project_id=project_id,\n",
    "        region=region,\n",
    "        master_image_uri=TORCH_TRAINER_IMAGE,\n",
    "        args=torch_args).set_display_name('Pytorch Model - AI Platform Training')\n",
    "    \n",
    "    train_xgb = mlengine_train_op(\n",
    "        project_id=project_id,\n",
    "        region=region,\n",
    "        master_image_uri=XGB_TRAINER_IMAGE,\n",
    "        args=xgb_args).set_display_name('XGBoost Model - AI Platform Training')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set environment variables for the different trainer image names as well as our bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TAG = 'latest'\n",
    "SCIKIT_TRAINER_IMAGE = 'gcr.io/{}/scikit_trainer_image:{}'.format(PROJECT_ID, TAG)\n",
    "TF_TRAINER_IMAGE = 'gcr.io/{}/tensorflow_trainer_image:{}'.format(PROJECT_ID, TAG)\n",
    "TORCH_TRAINER_IMAGE = 'gcr.io/{}/pytorch_trainer_image:{}'.format(PROJECT_ID, TAG)\n",
    "XGB_TRAINER_IMAGE = 'gcr.io/{}/xgboost_trainer_image:{}'.format(PROJECT_ID, TAG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env TF_TRAINER_IMAGE={TF_TRAINER_IMAGE}\n",
    "%env SCIKIT_TRAINER_IMAGE={SCIKIT_TRAINER_IMAGE}\n",
    "%env TORCH_TRAINER_IMAGE={TORCH_TRAINER_IMAGE}\n",
    "%env XGB_TRAINER_IMAGE={XGB_TRAINER_IMAGE}\n",
    "%env BUCKET={BUCKET}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile the Pipeline\n",
    "Compile the pipeline with the CLI compiler. This will save a census_training_pipeline.yaml file locally "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!dsl-compile --py pipeline/census_training_pipeline.py --output census_training_pipeline.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Take a look at the head of the yaml file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head census_training_pipeline.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set the command fields in the pipeline YAML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -i 's/\\\"command\\\": \\[\\]/\\\"command\\\": \\[python, -u, -m, kfp_component.launcher\\]/g' census_training_pipeline.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat census_training_pipeline.yaml | grep \"component.launcher\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see 6 lines in the output that were modified by the sed command."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy your KubeFlow Pipeline\n",
    "Now let's deploy the KubeFlow pipeline. Prior to the lab you should have spun up an AI Platform Pipelines Instance. In the AI Platform Pipeline UI click 'settings' on your pipeline and copy the endpoint. Paste the endpoint as the value for the string variable ENDPOINT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: Change ENDPOINT to the ENDPOINT for your AI Platform Pipelines Instance\n",
    "ENDPOINT = ''\n",
    "PIPELINE_NAME = 'census_trainer_multiple_models'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kfp --endpoint $ENDPOINT pipeline upload \\\n",
    "-p $PIPELINE_NAME \\\n",
    "./census_training_pipeline.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continuous Training: Create Pipeline Run, then Create Recurring Runs in the UI. \n",
    "Now that we deployed our KubeFlow pipeline, let's head over to the UI and launch a pipeline run. You can either click **Open Pipelines Dashboard** in the AI Platform Pipeline UI or run the Python cell below and copy/paste the output into the browser. Once in the UI, complete the following steps:\n",
    "* Select **Pipelines** on the left-hand navigation panel\n",
    "* Select **census_trainer_multiple_models** then click **Create Run**\n",
    "* In the **Experiment** field select **Default**\n",
    "* For **Run Type** select **One-Off**\n",
    "* Enter your **Project ID** and hit **Start**. You can now monitor your pipeline run in the UI (it will take about **10 minutes** to complete the run).\n",
    "\n",
    "\n",
    "**NOTE that your pipeline run may fail due to the bug in a BigQuery component that does not handle certain race conditions. If you observe the pipeline failure, please click on Retry on the top right in the KFP UI to re-run the failed steps.**\n",
    "\n",
    "\n",
    "Now let's set up a recurring run:\n",
    "* Select **Pipelines** then census_trainer_multiple_models and click **Create Run**\n",
    "* In the **Experiement** field select **Default**\n",
    "* For **Run Type** select **Recurring**\n",
    "* Configure the **Recurring Run** to start tomorrow at 5pm and run once weekly\n",
    "* Enter your **Project ID** and hit **Start**\n",
    "\n",
    "Now your pipeline will run weekly starting tomorrow, it's as easy as that!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"https://{ENDPOINT}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### End of Lab, Congrats!\n",
    "In this lab you created containerized training applications for TensorFlow, PyTorch, Scikit-learn, an XGBoost models. You then created a KubeFlow pipeline that used pre built components to create training/validation splits in BigQuery data, export that data as CSV files to GCS, and launch AI Platform Training Jobs with the four containerized training applications. Finally, you ran the pipeline from the UI and set up Continuous Training to re-run your pipeline once a week!"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-gpu.2-1.m55",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-1:m55"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
